// "Copyright [2021] <Copyright Shengkai Lin>"

package cluster.network_device;


import cluster.network_device.INic;
import roce;

module Machine {
    parameters:
        @controlplane;
        int numGpus = default(8);
        int numNvlink_port = default(4);
        int my_index;
        int nic_type = default(0);  // udp-roce = 0; ib = 1

    gates:
        inout out_port[numGpus];  // Data plane
        inout controller_port;  // Control plane

    submodules:
        cpu[numGpus+1] : Cpu {
            parameters:
                my_index = index;
                my_global_index = index + my_index*numGpus;
        }
        gpu[numGpus] : Gpu {
            parameters:
                numNvlink_port = numNvlink_port;
                my_global_index = index + my_index*numGpus;
                my_index = index;
        }
        nic[numGpus] : <default("EthernetNic")> like INic{
        // nic[numGpus] : <default("SimpleNic")> like INic{
            parameters:
                my_index = index;
                my_global_index = index + my_index*numGpus;
        }


    connections allowunconnected:
        // Only connect cpu0 to the controller directly.
        cpu[0].controller_port$o --> controller_port$o;
        cpu[0].controller_port$i <-- controller_port$i;

        for i=1..(numGpus) {
            cpu[i].controller_port <-->Pcie<--> cpu[0].cpu_port++;
            gpu[i-1].manager_port <-->Pcie<--> cpu[i].gpu_port;
            gpu[i-1].socket_port <-->Pcie<--> nic[i-1].gpu_port;
            nic[i-1].out_port <--> out_port++;
        }

        // The connections of NVlink
        // The sequence has been modified
        gpu[0].nvlink_port[1] <-->Nvlink<--> gpu[1].nvlink_port[3] if numNvlink_port==4;
        gpu[1].nvlink_port[1] <-->Nvlink<--> gpu[2].nvlink_port[3] if numNvlink_port==4;
        gpu[2].nvlink_port[1] <-->Nvlink<--> gpu[3].nvlink_port[3] if numNvlink_port==4;
        gpu[3].nvlink_port[1] <-->Nvlink<--> gpu[0].nvlink_port[3] if numNvlink_port==4;

        gpu[0].nvlink_port[2] <-->Nvlink<--> gpu[2].nvlink_port[2] if numNvlink_port==4;
        gpu[1].nvlink_port[2] <-->Nvlink<--> gpu[3].nvlink_port[2] if numNvlink_port==4;


        gpu[4].nvlink_port[1] <-->Nvlink<--> gpu[5].nvlink_port[3] if numNvlink_port==4;
        gpu[5].nvlink_port[1] <-->Nvlink<--> gpu[6].nvlink_port[3] if numNvlink_port==4;
        gpu[6].nvlink_port[1] <-->Nvlink<--> gpu[7].nvlink_port[3] if numNvlink_port==4;
        gpu[7].nvlink_port[1] <-->Nvlink<--> gpu[4].nvlink_port[3] if numNvlink_port==4;

        gpu[4].nvlink_port[2] <-->Nvlink<--> gpu[6].nvlink_port[2] if numNvlink_port==4;
        gpu[5].nvlink_port[2] <-->Nvlink<--> gpu[7].nvlink_port[2] if numNvlink_port==4;

        for i=0 .. 3 {
            gpu[i].nvlink_port[0] <-->Nvlink<--> gpu[i+4].nvlink_port[0] if numNvlink_port==4;
        }
        // gpu[0].nvlink_port[0] <-->Nvlink<--> gpu[7].nvlink_port[0] if numNvlink_port==4;
        // gpu[3].nvlink_port[0] <-->Nvlink<--> gpu[4].nvlink_port[0] if numNvlink_port==4;

}